#include "unp.h"

int main() {
	char buf[MAXLINE];
	int listenfd, servfd;
	int cli_fds[1024];
	int maxi;

	struct sockaddr_in serv_addr, cli_addr;
	socklen_t addr_len;

	if((listenfd = socket(AF_INET, SOCK_STREAM, 0)) == -1) 
		err_exit("socket listen fd error.");
	memset(&serv_addr, 0, sizeof(serv_addr));
	serv_addr.sin_family = AF_INET;
	serv_addr.sin_port = htons(9987);
	serv_addr.sin_addr.s_addr = htonl(INADDR_ANY);
	int opt_val = 1;
	if(setsockopt(listenfd, SOL_SOCKET, SO_REUSEADDR, &opt_val, sizeof(int)) == -1) 
		err_exit("set sock opt REUSEADDR error.");
	if(bind(listenfd, (struct sockaddr*)&serv_addr, sizeof(serv_addr)) != 0) 
		err_exit("bind serv addr error.");
	if(listen(listenfd, 1024) != 0)
		err_exit("listen error.");

	if((servfd = socket(AF_INET, SOCK_DGRAM, 0)) == -1)
		err_exit("socket udp fd error.");
	if(bind(servfd, (struct sockaddr*)&serv_addr, sizeof(serv_addr)) != 0)
		err_exit("bind udp serv addr error.");

	fd_set rfds;
	int nfds = listenfd > servfd ? listenfd : servfd;
	nfds+=1;
	FD_ZERO(&rfds);
	FD_SET(listenfd, &rfds);
	FD_SET(servfd, &rfds);

	maxi = 0;
	struct timeval timeout;
	timeout.tv_sec = 0;
	timeout.tv_usec = 10000;
	for(;;) {
		fd_set new_rfds = rfds;
		int nready = select(nfds, &new_rfds, NULL, NULL, &timeout);
		if(nready < 0) {
			if(errno == EINTR) continue;
			else err_exit("select error.");
		}
		if(nready == 0) continue;

		if(FD_ISSET(listenfd, &new_rfds)) {
			int cli_fd = accept(listenfd, NULL, NULL);
			if(cli_fd > 0) {
				int i;
				for(i = 0; i < maxi; ++i) {
					if(cli_fds[i] == -1) break;
				}
				if(i >= 1024) {
					err_msg("too many tcp connections.");
				} else {
					cli_fds[i] = cli_fd;
					FD_SET(cli_fd, &rfds);
					nfds = cli_fd > nfds -1 ? cli_fd + 1 : nfds;
					if(i == maxi) maxi++;
				}
			}
			if((--nready) <= 0) continue;
		}
		if(FD_ISSET(servfd, &new_rfds)) {
			addr_len = sizeof(cli_addr);
			int n = 0;
			if((n = recvfrom(servfd, buf, MAXLINE, 0, (struct sockaddr*)&cli_addr, &addr_len)) > 0) {
				sendto(servfd, buf, n, 0, (struct sockaddr*)&cli_addr, addr_len);
			}
			if((--nready) <= 0) continue;
		}

		for(int i = 0; i < maxi; ++i) {
			if(cli_fds[i] == -1) continue;
			if(! FD_ISSET(cli_fds[i], &new_rfds)) continue;
			int n;
			if((n = Read(cli_fds[i], buf, MAXLINE)) > 0) {
				Writen(cli_fds[i], buf, n);
			}
			if(n <= 0) {
				FD_CLR(cli_fds[i], &rfds);
				close(cli_fds[i]);
				cli_fds[i] = -1;
			}
			if((--nready) <= 0) break;
		}

	}
	return 0;
}
